// Author:   Tong Qin               qintonguav@gmail.com
// 	         Shaozu Cao 		    saozu.cao@connect.ust.hk

#ifndef LIO_AUTOFACTOR_HPP_
#define LIO_AUTOFACTOR_HPP_

#include <ceres/ceres.h>
#include <ceres/rotation.h>
#include <eigen3/Eigen/Dense>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl_conversions/pcl_conversions.h>
#include "utils/math_utils.h"

namespace lio{
	
struct LidarPlaneNormFactor
{

	LidarPlaneNormFactor(Eigen::Vector3d curr_point_, Eigen::Vector3d plane_unit_norm_,
						 double negative_OA_dot_norm_) : curr_point(curr_point_), plane_unit_norm(plane_unit_norm_),
														 negative_OA_dot_norm(negative_OA_dot_norm_) {}

	template <typename T>
	bool operator()(const T *transformTobeMapped, T *residual) const
	{
		Eigen::Matrix<T, 3, 1> cp{T(curr_point.x()), T(curr_point.y()), T(curr_point.z())};
		T x01 = cos(transformTobeMapped[2]) * cp.x()
				- sin(transformTobeMapped[2]) * cp.y();
		T y01 = sin(transformTobeMapped[2])* cp.x()
				+ cos(transformTobeMapped[2]) * cp.y();
		T z01 = cp.z();

		T x02 = x01;
		T y02 = cos(transformTobeMapped[0]) * y01 - sin(transformTobeMapped[0]) * z01;
		T z02 = sin(transformTobeMapped[0]) * y01 + cos(transformTobeMapped[0]) * z01;
		
		Eigen::Matrix<T, 3, 1> point_w;
		point_w.x() = cos(transformTobeMapped[1]) * x02 + sin(transformTobeMapped[1]) * z02
						+ transformTobeMapped[3];
		point_w.y() = y02 + transformTobeMapped[4];
		point_w.z() = -sin(transformTobeMapped[1]) * x02 + cos(transformTobeMapped[1]) * z02
						+ transformTobeMapped[5];	

		Eigen::Matrix<T, 3, 1> norm(T(plane_unit_norm.x()), T(plane_unit_norm.y()), T(plane_unit_norm.z()));
		residual[0] = norm.dot(point_w) + T(negative_OA_dot_norm);
		return true;
	}

	static ceres::CostFunction *Create(const Eigen::Vector3d curr_point_, const Eigen::Vector3d plane_unit_norm_,
									   const double negative_OA_dot_norm_)
	{
		return (new ceres::AutoDiffCostFunction<
				LidarPlaneNormFactor, 1, 6>(
			new LidarPlaneNormFactor(curr_point_, plane_unit_norm_, negative_OA_dot_norm_)));
	}

	Eigen::Vector3d curr_point;
	Eigen::Vector3d plane_unit_norm;
	double negative_OA_dot_norm;
};

struct LidarEdgeFactorPose
{
	LidarEdgeFactorPose(Eigen::Vector3d curr_point_, Eigen::Vector3d last_point_a_,
					Eigen::Vector3d last_point_b_, double s_)
		: curr_point(curr_point_), last_point_a(last_point_a_), last_point_b(last_point_b_), s(s_) {}

	template <typename T>
	bool operator()(const T *transformTobeMapped, T *residual) const
	{

		Eigen::Matrix<T, 3, 1> cp{T(curr_point.x()), T(curr_point.y()), T(curr_point.z())};
		Eigen::Matrix<T, 3, 1> lpa{T(last_point_a.x()), T(last_point_a.y()), T(last_point_a.z())};
		Eigen::Matrix<T, 3, 1> lpb{T(last_point_b.x()), T(last_point_b.y()), T(last_point_b.z())};
		
		T x01 = cos(transformTobeMapped[2]) * cp.x()
				- sin(transformTobeMapped[2]) * cp.y();
		T y01 = sin(transformTobeMapped[2])* cp.x()
				+ cos(transformTobeMapped[2]) * cp.y();
		T z01 = cp.z();

		T x02 = x01;
		T y02 = cos(transformTobeMapped[0]) * y01 - sin(transformTobeMapped[0]) * z01;
		T z02 = sin(transformTobeMapped[0]) * y01 + cos(transformTobeMapped[0]) * z01;
		
		Eigen::Matrix<T, 3, 1> lp;
		lp.x() = cos(transformTobeMapped[1]) * x02 + sin(transformTobeMapped[1]) * z02
						+ transformTobeMapped[3];
		lp.y() = y02 + transformTobeMapped[4];
		lp.z() = -sin(transformTobeMapped[1]) * x02 + cos(transformTobeMapped[1]) * z02
						+ transformTobeMapped[5];	

		Eigen::Matrix<T, 3, 1> nu = (lp - lpa).cross(lp - lpb);
		Eigen::Matrix<T, 3, 1> de = lpa - lpb;

		residual[0] = nu.x() / de.norm();
		residual[1] = nu.y() / de.norm();
		residual[2] = nu.z() / de.norm();

		return true;
	}

	static ceres::CostFunction *Create(const Eigen::Vector3d curr_point_, const Eigen::Vector3d last_point_a_,
									   const Eigen::Vector3d last_point_b_, const double s_)
	{
		return (new ceres::AutoDiffCostFunction<
				LidarEdgeFactorPose, 3, 6>(
			new LidarEdgeFactorPose(curr_point_, last_point_a_, last_point_b_, s_)));
	}

	Eigen::Vector3d curr_point, last_point_a, last_point_b;
	double s;
};

struct LioLidarPlaneNormFactor
{

	LioLidarPlaneNormFactor(Eigen::Vector3d curr_point_, Eigen::Vector3d plane_unit_norm_,
						 double negative_OA_dot_norm_) : curr_point(curr_point_), plane_unit_norm(plane_unit_norm_),
														 negative_OA_dot_norm(negative_OA_dot_norm_) {}

	template <typename T>
	bool operator()(const T *qandt, T *residual) const
	{
		Eigen::Quaternion<T> q_w_curr{qandt[3], qandt[0], qandt[1], qandt[2]};
		Eigen::Matrix<T, 3, 1> t_w_curr{qandt[4], qandt[5], qandt[6]};
		Eigen::Matrix<T, 3, 1> cp{T(curr_point.x()), T(curr_point.y()), T(curr_point.z())};
		Eigen::Matrix<T, 3, 1> point_w;
		point_w = q_w_curr * cp + t_w_curr;

		Eigen::Matrix<T, 3, 1> norm(T(plane_unit_norm.x()), T(plane_unit_norm.y()), T(plane_unit_norm.z()));
		residual[0] = norm.dot(point_w) + T(negative_OA_dot_norm);
		return true;
	}

	static ceres::CostFunction *Create(const Eigen::Vector3d curr_point_, const Eigen::Vector3d plane_unit_norm_,
									   const double negative_OA_dot_norm_)
	{
		return (new ceres::AutoDiffCostFunction<
				LioLidarPlaneNormFactor, 1, 7>(
			new LioLidarPlaneNormFactor(curr_point_, plane_unit_norm_, negative_OA_dot_norm_)));
	}

	Eigen::Vector3d curr_point;
	Eigen::Vector3d plane_unit_norm;
	double negative_OA_dot_norm;
};

struct LioLidarEdgeFactorPose
{
	LioLidarEdgeFactorPose(Eigen::Vector3d curr_point_, Eigen::Vector3d last_point_a_,
					Eigen::Vector3d last_point_b_, double s_)
		: curr_point(curr_point_), last_point_a(last_point_a_), last_point_b(last_point_b_), s(s_) {}

	template <typename T>
	bool operator()(const T *qandt, T *residual) const
	{

		Eigen::Matrix<T, 3, 1> cp{T(curr_point.x()), T(curr_point.y()), T(curr_point.z())};
		Eigen::Matrix<T, 3, 1> lpa{T(last_point_a.x()), T(last_point_a.y()), T(last_point_a.z())};
		Eigen::Matrix<T, 3, 1> lpb{T(last_point_b.x()), T(last_point_b.y()), T(last_point_b.z())};

		//Eigen::Quaternion<T> q_last_curr{q[3], T(s) * q[0], T(s) * q[1], T(s) * q[2]};
		Eigen::Quaternion<T> q_last_curr{qandt[3], qandt[0], qandt[1], qandt[2]};
		Eigen::Quaternion<T> q_identity{T(1), T(0), T(0), T(0)};
		q_last_curr = q_identity.slerp(T(s), q_last_curr);
		Eigen::Matrix<T, 3, 1> t_last_curr{T(s) * qandt[4], T(s) * qandt[5], T(s) * qandt[6]};

		Eigen::Matrix<T, 3, 1> lp;
		lp = q_last_curr * cp + t_last_curr;

		Eigen::Matrix<T, 3, 1> nu = (lp - lpa).cross(lp - lpb);
		Eigen::Matrix<T, 3, 1> de = lpa - lpb;

		residual[0] = nu.x() / de.norm();
		residual[1] = nu.y() / de.norm();
		residual[2] = nu.z() / de.norm();

		return true;
	}

	static ceres::CostFunction *Create(const Eigen::Vector3d curr_point_, const Eigen::Vector3d last_point_a_,
									   const Eigen::Vector3d last_point_b_, const double s_)
	{
		return (new ceres::AutoDiffCostFunction<
				LioLidarEdgeFactorPose, 3, 7>(
			new LioLidarEdgeFactorPose(curr_point_, last_point_a_, last_point_b_, s_)));
	}

	Eigen::Vector3d curr_point, last_point_a, last_point_b;
	double s;
};


// 手动求导，没有旋转约束 2021/12/22
using namespace mathutils;

class AnalyticLidarPlaneNormFactor : public ceres::SizedCostFunction<1, 7> {

public:
	AnalyticLidarPlaneNormFactor(Eigen::Vector3d curr_point_, Eigen::Vector3d plane_unit_norm_,
							double negative_OA_dot_norm_) : curr_point(curr_point_), plane_unit_norm(plane_unit_norm_),
															negative_OA_dot_norm(negative_OA_dot_norm_) {}

	virtual bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const {
		Eigen::Quaterniond q_w_curr{parameters[0][3], parameters[0][0], parameters[0][1], parameters[0][2]};
		Eigen::Matrix<double, 3, 1> t_w_curr{parameters[0][4], parameters[0][5], parameters[0][6]};

		Eigen::Matrix<double, 3, 1> point_w;
		point_w = q_w_curr * curr_point + t_w_curr;

		residuals[0] = plane_unit_norm.dot(point_w) + negative_OA_dot_norm;

		if (jacobians) {
			// Eigen::Matrix3d Rwc = q_w_curr.toRotationMatrix();
			if (jacobians[0]) {
				Eigen::Map<Eigen::Matrix<double, 1, 7, Eigen::RowMajor> > jacobian_pq(jacobians[0]);
				jacobian_pq.setZero();

				// jacobian_pq.block<1, 3>(0, 0) = -plane_unit_norm.transpose() * (Rwc * SkewSymmetric(curr_point)) * Rwc.transpose();
				jacobian_pq.block<1, 3>(0, 0) = -plane_unit_norm.transpose() * (q_w_curr * SkewSymmetric(curr_point)) * q_w_curr.inverse().toRotationMatrix();
				jacobian_pq.block<1, 3>(0, 3) = plane_unit_norm.transpose();
			}    
		}

		return true;
	}

	void Check(double **parameters);

	Eigen::Vector3d curr_point;
	Eigen::Vector3d plane_unit_norm;
	double negative_OA_dot_norm;

	EIGEN_MAKE_ALIGNED_OPERATOR_NEW

};

class AnalyticLidarEdgeFactor : public ceres::SizedCostFunction<3, 7> {

public:
	AnalyticLidarEdgeFactor(Eigen::Vector3d curr_point_, Eigen::Vector3d last_point_a_,
					Eigen::Vector3d last_point_b_)
		: curr_point(curr_point_), last_point_a(last_point_a_), last_point_b(last_point_b_) {}

	virtual bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const {
		Eigen::Quaterniond q_w_curr{parameters[0][3], parameters[0][0], parameters[0][1], parameters[0][2]};
		Eigen::Matrix<double, 3, 1> t_w_curr{parameters[0][4], parameters[0][5], parameters[0][6]};

		Eigen::Matrix<double, 3, 1> lp;
		lp = q_w_curr * curr_point + t_w_curr;

		Eigen::Matrix<double, 3, 1> nu = (lp - last_point_a).cross(lp - last_point_b);
		Eigen::Matrix<double, 3, 1> de = last_point_a - last_point_b;

		residuals[0] = nu.x() / de.norm();
		residuals[1] = nu.y() / de.norm();
		residuals[2] = nu.z() / de.norm();

		if (jacobians) {
			Eigen::Matrix3d Rwc = q_w_curr.toRotationMatrix();
			Eigen::Matrix3d dre_dx = SkewSymmetric(last_point_b - last_point_a) / de.norm();
			if (jacobians[0]) {
				Eigen::Map<Eigen::Matrix<double, 3, 7, Eigen::RowMajor> > jacobian_pq(jacobians[0]);
				jacobian_pq.setZero();

				jacobian_pq.block<3, 3>(0, 0) = -dre_dx * (Rwc * SkewSymmetric(curr_point)) * Rwc.transpose();
				jacobian_pq.block<3, 3>(0, 3) = dre_dx;
			}    
		}

		return true;
	}

	void Check(double **parameters);

	Eigen::Vector3d curr_point, last_point_a, last_point_b;

	EIGEN_MAKE_ALIGNED_OPERATOR_NEW

};

// class AnalyticLidarEdgeFactor : public ceres::SizedCostFunction<1, 7> {

// public:
// 	AnalyticLidarEdgeFactor(Eigen::Vector3d curr_point_, Eigen::Vector3d last_point_a_,
// 					Eigen::Vector3d last_point_b_)
// 		: curr_point(curr_point_), last_point_a(last_point_a_), last_point_b(last_point_b_) {}

// 	virtual bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const {
// 		Eigen::Quaterniond q_w_curr{parameters[0][3], parameters[0][0], parameters[0][1], parameters[0][2]};
// 		Eigen::Matrix<double, 3, 1> t_w_curr{parameters[0][4], parameters[0][5], parameters[0][6]};

// 		Eigen::Matrix<double, 3, 1> lp;
// 		lp = q_w_curr * curr_point + t_w_curr;

// 		Eigen::Matrix<double, 3, 1> nu = (lp - last_point_a).cross(lp - last_point_b);
// 		Eigen::Matrix<double, 3, 1> de = last_point_a - last_point_b;

// 		residuals[0] = nu.norm() / de.norm();

// 		if (jacobians) {
// 			// Eigen::Matrix3d Rwc = q_w_curr.toRotationMatrix();
// 			Eigen::Vector3d w = (last_point_a - last_point_b).cross(nu);
// 			w.normalized();
// 			if (jacobians[0]) {
// 				Eigen::Map<Eigen::Matrix<double, 1, 7, Eigen::RowMajor> > jacobian_pq(jacobians[0]);
// 				jacobian_pq.setZero();

// 				// jacobian_pq.block<1, 3>(0, 0) = -w.transpose() * (Rwc * SkewSymmetric(curr_point)) * Rwc.transpose();
// 				jacobian_pq.block<1, 3>(0, 0) = -w.transpose() * (q_w_curr * SkewSymmetric(curr_point)) * q_w_curr.inverse().toRotationMatrix();
// 				jacobian_pq.block<1, 3>(0, 3) = w.transpose();
// 			}    
// 		}

// 		return true;
// 	}

// 	void Check(double **parameters);

// 	Eigen::Vector3d curr_point, last_point_a, last_point_b;

// 	EIGEN_MAKE_ALIGNED_OPERATOR_NEW

// };


}

#endif
